from absl import app, flags

import datetime
import os
import glob
import numpy
import pandas
import logging

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS

import coin.strategy.mm.tool.plot_util as putil


def from_fill_csv_group(fill_csv_files):
  dfs = []
  for fill_csv_file in fill_csv_files:
    try:
      df = pandas.read_csv(fill_csv_file)
      df['balance_diff'] = numpy.nan_to_num(df['balance'].diff())
      dfs.append(df)
    except Exception:
      pass
  if len(dfs) == 0:
    return None
  return pandas.concat(dfs, axis=0).sort_values('time')


def png_from_fill_csv(df, out_png_file):
  df['balance2'][df['time'].diff() > 6000e9] = numpy.nan
  plt.plot(pandas.DatetimeIndex(df['time']), df['balance2'], linewidth=0.5)
  # plt.plot(pandas.DatetimeIndex(df['time']), df['balance2'], 'r.', markersize=0.2)
  plt.title(out_png_file.split("/")[-1])
  plt.xticks(rotation=10)
  plt.tight_layout()
  plt.savefig(out_png_file)
  plt.close()


def extract_group_from_filename(csv_file):
  csv_file_splits = csv_file.split(".")
  dt_col = None
  for i, csv_file_split in enumerate(csv_file_splits):
    try:
      assert dt_col is None
      datetime.datetime.strptime(csv_file_split, '%Y%m%d')
      dt_col = i
    except Exception:
      pass
  assert dt_col is not None
  csv_file_splits[dt_col] = '*'
  return ".".join(csv_file_splits)


def main(_):
  filepattern = os.path.expanduser(FLAGS.fill_csv_file)
  logging.info(filepattern)
  group_globs = set()
  for csv_file in glob.glob(filepattern):
    group_glob = extract_group_from_filename(csv_file)
    group_globs.add(group_glob)

  for group_glob in group_globs:
    try:
      df = from_fill_csv_group(glob.glob(group_glob))
      df['balance2'] = df['balance_diff'].cumsum()
      os.makedirs("oe_stat_backtest", exist_ok=True)
      out_filename = os.path.join("oe_stat_backtest", "pnl_" + group_glob.split("/")[-1] + ".png")
      png_from_fill_csv(df, out_filename)
      logging.info(out_filename)
    except Exception:
      pass


if __name__ == '__main__':
  flags.DEFINE_string('fill_csv_file', None, 'fill_csv_file mode if given')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
